from unittest.mock import patch, MagicMock
from library_test.base_test import BaseTest


class BaseTestCheckUtil(BaseTest):

    @classmethod
    def get_module_path(cls):
        return "ascend_deployer.module_utils.check_utils"

    @classmethod
    def setUpClass(cls) -> None:
        super().setUpClass()
        from ascend_deployer.module_utils.check_utils import CheckUtil
        cls.check_util = CheckUtil()


class TestCheckUtils(BaseTestCheckUtil):

    def test_get_card(self):
        patcher = patch(self.get_module_path() + ".get_npu_info")
        get_npu_info_mocker = patcher.start()
        get_npu_info_mocker.return_value = {"scene": "infer"}
        self.assertEqual("--", self.check_util.get_card())

        get_npu_info_mocker.return_value = {"scene": "train"}
        self.assertEqual("910", self.check_util.get_card())

        get_npu_info_mocker.return_value = {"scene": "a300i"}
        self.assertEqual("310p", self.check_util.get_card())

        get_npu_info_mocker.return_value = {"scene": "a910b"}
        self.assertEqual("910b", self.check_util.get_card())

        get_npu_info_mocker.return_value = {"scene": "a910_93"}
        self.assertEqual("910_93", self.check_util.get_card())

    def test_run_cmd(self):
        patcher = patch(self.get_module_path() + ".subprocess.Popen")
        subprocess_mocker = patcher.start()
        res_mocker = MagicMock()

        # success_code does not match the returncode
        res_mocker.communicate.return_value = ("failed", "")
        res_mocker.returncode = 3
        subprocess_mocker.return_value = res_mocker

        with self.assertRaises(Exception) as context:
            self.check_util.run_cmd("ps -a", self.check_util.GREP_RETURN_CODE)
        self.assertEqual("call cmd ps -a failed, reason: failed", str(context.exception))

        # matched returncode
        res_mocker.communicate.return_value = ("success", "")
        res_mocker.returncode = 1
        subprocess_mocker.return_value = res_mocker
        self.assertEqual("success", self.check_util.run_cmd("ps -a", self.check_util.GREP_RETURN_CODE))

        # no success_code and returncode is not 0
        res_mocker.communicate.return_value = ("failed", "")
        res_mocker.returncode = 1
        subprocess_mocker.return_value = res_mocker
        with self.assertRaises(Exception) as context:
            self.check_util.run_cmd("ps -a")
        self.assertEqual("call cmd ps -a failed, reason: failed", str(context.exception))

        # no success_code and returncode is 0
        res_mocker.communicate.return_value = ("success", "")
        res_mocker.returncode = 0
        subprocess_mocker.return_value = res_mocker
        self.assertEqual("success", self.check_util.run_cmd("ps -a"))

    def test_record_error(self):
        error_messages = []
        msg = "test1"
        self.check_util.record_error(msg, error_messages)
        self.assertEqual([msg], error_messages)
        msg = "test2"
        self.check_util.record_error(msg, error_messages)
        self.assertEqual(["test1", "test2"], error_messages)
        # duplicated msg
        self.check_util.record_error(msg, error_messages)
        self.assertEqual(["test1", "test2"], error_messages)

    def test_find_file(self):
        patcher = patch(self.get_module_path() + ".os.walk")
        walk_mocker = patcher.start()
        walk_mocker.return_value = [
            ("/test1", ("subtest1", "subtest2"), ("s1.txt", "s2.txt")),
            ("/test2", (), ("1.txt", )),
            ("/test3", (), ("2.txt", ))
        ]
        self.assertEqual(None, self.check_util.find_file("resource", "file"))
        self.assertEqual("1.txt", self.check_util.find_file("resource", "1.txt"))
